package com.gitee.wsl.coroutines.executor

import com.gitee.wsl.concurrent.api.SuspendRunnable
import com.gitee.wsl.concurrent.api.SuspendRunnableWithResult
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.channels.BufferOverflow
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.consumeAsFlow
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.launch
import kotlin.coroutines.CoroutineContext


/**
 * A coroutine-based utility that coalesces multiple trigger calls into at most one active execution
 * of a suspending [task]. It runs the task in a dedicated coroutine and ensures that:
 *
 * - Only one instance of [task] runs at a time.
 * - If [schedule] is called while the task is running, the task is scheduled to run once more
 *   after the current execution finishes.
 * - Repeated calls to [schedule] while the task is running are coalesced into a single extra run.
 *
 * This is useful for scenarios like batching updates or refreshing state from external sources,
 * where overlapping invocations should not stack up but a final follow-up call should be ensured.
 *
 * Example:
 * ```
 * val runner = CoalescingTaskRunner(context = Dispatchers.Default) {
 *     println("Running update")
 *     delay(1000)
 * }
 *
 * repeat(10) { runner.schedule() } // Will run at most twice
 *
 * runner.close()
 * ```
 *
 * @param context Coroutine context to run the task (defaults to [Dispatchers.Default]).
 * @param errorHandler Optional callback for handling exceptions thrown by the task.
 * @param task The suspending function to execute. Only one instance runs at a time.
 */
class CoalescingTaskRunner(
    context: CoroutineContext = Dispatchers.Default,
    private val errorHandler: (Throwable) -> Unit = {},
    private val task: SuspendRunnable,
) {
    private val scope: CoroutineScope = CoroutineScope(context + SupervisorJob())
    private val triggerChannel = Channel<Unit>(
        capacity = 1,
        onBufferOverflow = BufferOverflow.DROP_OLDEST,
    )
    private val runner = scope.launch {
        triggerChannel.consumeAsFlow().collect {
            try {
                task()
            } catch (t: Throwable) {
                runCatching { errorHandler(t) }
            }
        }
    }

     fun schedule() {
        triggerChannel.trySend(Unit)
    }

     suspend fun close() {
        triggerChannel.close()
        runner.join()
        scope.cancel()
    }

     suspend fun <T> use(block: suspend (CoalescingTaskRunner) -> T): T = try {
        block(this)
    } finally {
        close()
    }
}


/**
 * A coroutine-based utility that coalesces multiple invocations of a suspending task
 * and ensures only the latest scheduled invocation is executed, returning its result.
 *
 * This is useful in scenarios where a task may be triggered frequently (e.g., UI events,
 * network polling), but only the most recent execution and its result are relevant.
 *
 * Example use case: auto-refreshing data on user input where intermediate results can be discarded.
 *
 * The task is executed on a dedicated coroutine scope with optional error handling.
 * If a task is currently running and a new one is scheduled, the new task supersedes
 * previous pending triggers but waits for its result.
 *
 * @param context Coroutine context to run the task (defaults to [Dispatchers.Default]).
 * @param task The suspending function to execute. Only one instance runs at a time.
 * @param T The result type of the task.
 */
class CoalescingTaskRunnerWithResult<T>(
    context: CoroutineContext = Dispatchers.Default,
    private val task: SuspendRunnableWithResult<T>,
) {
    private data class State<T>(
        val idRunning: Long?,
        val idLatest: Long,
        val value: T?,
    )

    private val result = MutableStateFlow<State<T>>(State(idRunning = null, idLatest = -1, value = null))
    private val scope: CoroutineScope = CoroutineScope(context + SupervisorJob())
    private val triggerChannel = Channel<Unit>(
        capacity = 1,
        onBufferOverflow = BufferOverflow.DROP_OLDEST,
    )
    private val runner = scope.launch {
        triggerChannel.consumeAsFlow().collect {
            val state = result.first()
            val idRunning = state.idLatest + 1
            result.value = state.copy(idRunning = idRunning)
            val value = task()
            result.value = State(idRunning = null, idLatest = idRunning, value = value)
        }
    }

     suspend fun schedule(): T {
        val latestId = result.first().let { it.idRunning ?: it.idLatest }
        triggerChannel.send(Unit)
        return result.first { it.idLatest > latestId }.let {
            @Suppress("UNCHECKED_CAST")
            it.value as T
        }
    }

     suspend fun close() {
        triggerChannel.close()
        runner.join()
        scope.cancel()
    }

     suspend fun <R> use(block: suspend (CoalescingTaskRunnerWithResult<T>) -> R): R = try {
        block(this)
    } finally {
        close()
    }
}